package pool

import (
	"sync"
	"sync/atomic"
)

var (
	bufPool64k sync.Pool
	bufPool32k sync.Pool
	bufPool16k sync.Pool
	bufPool8k  sync.Pool
	bufPool4k  sync.Pool
	bufPool2k  sync.Pool
	bufPool1k  sync.Pool
	bufPool512 sync.Pool
	bufPool256 sync.Pool
	bufPool    sync.Pool

	// 性能统计
	allocCount uint64
	hitCount   uint64
)

// 初始化缓冲池
func init() {
	bufPool64k = sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, 64*1024)
		},
	}
	bufPool32k = sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, 32*1024)
		},
	}
	bufPool16k = sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, 16*1024)
		},
	}
	bufPool8k = sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, 8*1024)
		},
	}
	bufPool4k = sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, 4*1024)
		},
	}
	bufPool2k = sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, 2*1024)
		},
	}
	bufPool1k = sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, 1*1024)
		},
	}
	bufPool512 = sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, 512)
		},
	}
	bufPool256 = sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, 256)
		},
	}
	bufPool = sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, 128)
		},
	}
}

// 优化的Malloc函数
func Malloc(size int) []byte {
	atomic.AddUint64(&allocCount, 1)

	var x any
	switch {
	case size >= 64*1024:
		x = bufPool64k.Get()
	case size >= 32*1024:
		x = bufPool32k.Get()
	case size >= 16*1024:
		x = bufPool16k.Get()
	case size >= 8*1024:
		x = bufPool8k.Get()
	case size >= 4*1024:
		x = bufPool4k.Get()
	case size >= 2*1024:
		x = bufPool2k.Get()
	case size >= 1*1024:
		x = bufPool1k.Get()
	case size >= 512:
		x = bufPool512.Get()
	case size >= 256:
		x = bufPool256.Get()
	default:
		x = bufPool.Get()
	}

	if x == nil {
		return make([]byte, size)
	}

	buf := x.([]byte)
	if cap(buf) < size {
		return make([]byte, size)
	}

	atomic.AddUint64(&hitCount, 1)
	return buf[:size]
}

// 优化的Free函数
func Free(buf []byte) {
	if cap(buf) == 0 {
		return
	}

	// 重置切片长度以避免内存泄漏
	buf = buf[:0]

	size := cap(buf)
	switch {
	case size >= 64*1024:
		bufPool64k.Put(buf)
	case size >= 32*1024:
		bufPool32k.Put(buf)
	case size >= 16*1024:
		bufPool16k.Put(buf)
	case size >= 8*1024:
		bufPool8k.Put(buf)
	case size >= 4*1024:
		bufPool4k.Put(buf)
	case size >= 2*1024:
		bufPool2k.Put(buf)
	case size >= 1*1024:
		bufPool1k.Put(buf)
	case size >= 512:
		bufPool512.Put(buf)
	case size >= 256:
		bufPool256.Put(buf)
	default:
		bufPool.Put(buf)
	}
}

// 获取性能统计
func GetStats() (uint64, uint64, float64) {
	allocs := atomic.LoadUint64(&allocCount)
	hits := atomic.LoadUint64(&hitCount)
	var hitRate float64
	if allocs > 0 {
		hitRate = float64(hits) / float64(allocs) * 100
	}
	return allocs, hits, hitRate
}

// 重置统计
func ResetStats() {
	atomic.StoreUint64(&allocCount, 0)
	atomic.StoreUint64(&hitCount, 0)
}
